from fontTools.misc.cython import returns
# 需要改变
from osc02.utils.random_util import RandomUtil

random_util = RandomUtil(seed=1)
floor_durations = random_util.get_sf_erection_durations()
std_duration = floor_durations[2:-1]

print(floor_durations, len(floor_durations))
# print(std_duration, len(std_duration))

floor_durations = random_util.get_sf_erection_durations()
std_duration = floor_durations[2:-1]

print(floor_durations, len(floor_durations))
# print(std_duration, len(std_duration))

exit()

import numpy as np
import scipy
from scipy.stats import (triang, norm, expon, poisson,
                         kstest, chisquare, anderson)
import matplotlib.pyplot as plt
from typing import Union, List, Dict, Optional


class DistributionTester:
    """
    分布检验类，支持多种概率分布的拟合检验

    支持的分布类型:
    - 'triangular': 三角分布 - 需要参数: left(最小值), mode(众数), right(最大值)
    - 'normal': 正态分布 - 需要参数: mean(均值), std(标准差)
    - 'exponential': 指数分布 - 需要参数: scale(尺度参数，等于1/λ)
    - 'poisson': 泊松分布(离散) - 需要参数: mu(均值/方差)

    支持的检验方法:
    - 'ks': Kolmogorov-Smirnov检验(适用于连续分布)
    - 'anderson': Anderson-Darling检验(适用于正态和指数分布)
    - 'chi2': 卡方检验(适用于离散分布)
    """

    def __init__(self):
        """
        初始化分布检验器，设置支持的分布类型和matplotlib字体配置
        """
        # 支持的分布类型及其参数配置
        self.supported_dists = {
            'triangular': {
                'params': ['left', 'mode', 'right'],  # 三角分布参数列表
                'continuous': True  # 是否为连续分布
            },
            'normal': {
                'params': ['mean', 'std'],
                'continuous': True
            },
            'exponential': {
                'params': ['scale'],
                'continuous': True
            },
            'poisson': {
                'params': ['mu'],
                'continuous': False  # 离散分布
            }
        }

        # 配置matplotlib中文字体和符号显示
        # 字体回退顺序：SimHei(黑体) -> Microsoft YaHei(雅黑) -> 系统默认无衬线字体
        plt.rcParams["font.family"] = ['SimHei', "Microsoft YaHei", "sans-serif"]
        # 解决负号(-)显示为方框的问题
        plt.rcParams['axes.unicode_minus'] = False
        np.random.seed
        print('scipy.__version__ = ', scipy.__version__)
        print('np.__version__ = ', np.__version__)

    def test_distribution(
            self,
            samples: Union[List[float], np.ndarray],
            dist_type: str,
            dist_params: Dict[str, float],
            alpha: float = 0.05,
            plot: bool = True,
            test_type: str = 'ks'
    ) -> Dict:
        """
        检验样本数据是否符合指定的概率分布

        参数:
            samples: 待检验的样本数据，可以是列表或numpy数组
            dist_type: 要检验的分布类型，必须是supported_dists中定义的
            dist_params: 分布参数字典，必须包含该分布所需的所有参数
            alpha: 显著性水平(默认0.05)，用于假设检验
            plot: 是否绘制样本与理论分布的比较图(默认True)
            test_type: 使用的检验方法('ks'/'anderson'/'chi2')，或自动选择

        返回:
            包含检验结果的字典，包含以下键:
            - 'statistic': 检验统计量
            - 'p_value': p值(如果检验提供)
            - 'result': 检验结果(True/False)
            - 'message': 可读的结果描述
            - 'alpha': 使用的显著性水平
            - 'sample_size': 样本量
            - 'distribution': 检验的分布类型
            - 'distribution_params': 使用的分布参数
            - 'test_type': 使用的检验方法
        """
        # 参数校验：检查分布类型是否支持
        if dist_type not in self.supported_dists:
            raise ValueError(f"不支持的分布类型。可选: {list(self.supported_dists.keys())}")

        # 检查是否提供了该分布所需的所有参数
        required_params = self.supported_dists[dist_type]['params']
        for param in required_params:
            if param not in dist_params:
                raise ValueError(f"缺失必要参数: {param}")

        # 将输入数据转换为numpy数组
        samples = np.asarray(samples)

        # 样本量检查，小样本可能导致检验不可靠
        if len(samples) < 20:
            print("警告: 样本量较小，检验结果可能不可靠")

        # 根据分布类型选择合适的检验方法
        is_continuous = self.supported_dists[dist_type]['continuous']

        # 执行选定的检验方法
        if test_type == 'ks' and is_continuous:
            test_result = self._ks_test(samples, dist_type, dist_params)
        elif test_type == 'anderson' and is_continuous:
            test_result = self._anderson_test(samples, dist_type, dist_params, alpha)
        elif test_type == 'chi2' and not is_continuous:
            test_result = self._chi2_test(samples, dist_type, dist_params)
        else:
            raise ValueError(f"不支持的检验组合: {dist_type} with {test_type}")

        # 构建完整的检验结果字典
        test_result.update({
            'alpha': alpha,  # 显著性水平
            'sample_size': len(samples),  # 样本量
            'distribution': dist_type,  # 检验的分布类型
            'distribution_params': dist_params,  # 使用的分布参数
            'test_type': test_type,  # 使用的检验方法
            # 判断是否拒绝原假设(p值>alpha则不拒绝)
            'result': test_result['p_value'] > alpha if 'p_value' in test_result else None,
            # 生成可读的结果描述
            'message': self._get_result_message(test_result, alpha)
        })

        # 如果需要，绘制样本与理论分布的比较图
        if plot:
            self._plot_comparison(samples, dist_type, dist_params, test_result)

        return test_result

    def _ks_test(self, samples, dist_type, dist_params):
        if dist_type == 'triangular':
            left, mode, right = dist_params['left'], dist_params['mode'], dist_params['right']
            c = (mode - left) / (right - left)

            # 自定义CDF函数避免参数传递问题
            def triangular_cdf(x):
                x_scaled = (np.clip(x, left, right) - left) / (right - left)
                return triang._cdf(x_scaled, c)

            statistic, p_value = kstest(samples, triangular_cdf)
        elif dist_type == 'normal':
            statistic, p_value = kstest(samples, 'norm',
                                        args=(dist_params['mean'], dist_params['std']))
        elif dist_type == 'exponential':
            statistic, p_value = kstest(samples, 'expon',
                                        args=(0, dist_params['scale']))
        else:
            raise ValueError("KS检验不支持此分布类型")
        return {'statistic': statistic, 'p_value': p_value}

    def _anderson_test(self, samples, dist_type, dist_params, alpha):
        """
        Anderson-Darling检验(适用于正态和指数分布)
        比KS检验对尾部差异更敏感

        参数:
            samples: 样本数据
            dist_type: 分布类型
            dist_params: 分布参数(虽然AD检验不直接使用这些参数)
            alpha: 显著性水平

        返回:
            包含以下内容的字典:
            - 'statistic': AD统计量
            - 'critical_value': 临界值
            - 'alpha': 实际使用的显著性水平
            - 'result': 检验结果(True/False)
        """
        # 分布类型映射到scipy的命名
        dist_map = {
            'normal': 'norm',
            'exponential': 'expon',
            'triangular': None  # Anderson检验不支持三角分布
        }

        if dist_type not in dist_map or dist_map[dist_type] is None:
            raise ValueError("Anderson检验不支持此分布类型")

        # 执行Anderson-Darling检验
        result = anderson(samples, dist_map[dist_type])

        # 查找与指定alpha对应的临界值
        critical_values = result.critical_values
        significance_level = result.significance_level  # scipy中这是百分比值(如5表示5%)
        idx = np.where(significance_level == alpha * 100)[0]

        if len(idx) == 0:
            # 如果没有精确匹配的alpha，使用最接近的值
            closest_idx = np.argmin(np.abs(significance_level - alpha * 100))
            used_alpha = significance_level[closest_idx] / 100
            print(f"警告: 使用最接近的显著性水平 {used_alpha} 代替 {alpha}")
        else:
            used_alpha = alpha
            closest_idx = idx[0]

        return {
            'statistic': result.statistic,  # AD统计量
            'critical_value': critical_values[closest_idx],  # 临界值
            'alpha': used_alpha,  # 实际使用的alpha
            # 比较统计量与临界值
            'result': result.statistic < critical_values[closest_idx]
        }

    def _chi2_test(self, samples, dist_type, dist_params):
        """
        卡方拟合优度检验(适用于离散分布)
        比较观察频数与理论频数的差异

        参数:
            samples: 样本数据
            dist_type: 分布类型(目前仅支持泊松)
            dist_params: 分布参数

        返回:
            包含以下内容的字典:
            - 'statistic': 卡方统计量
            - 'p_value': p值
            - 'df': 自由度
            - 'observed': 观察频数
            - 'expected': 理论频数
        """
        if dist_type != 'poisson':
            raise ValueError("卡方检验当前仅支持泊松分布")

        # 计算每个唯一值的观察频数
        unique, counts = np.unique(samples, return_counts=True)
        observed = counts

        # 计算理论频数 = 概率质量函数(PMF) × 总样本数
        expected = poisson.pmf(unique, dist_params['mu']) * len(samples)

        # 合并小频数组(卡方检验要求每组期望频数>=5)
        while True:
            # 归一化期望频数总和，使其与观察频数总和一致
            expected = expected * sum(observed) / sum(expected)

            # 检查是否所有期望频数>=5 或只剩2组
            if np.all(expected >= 5) or len(expected) <= 2:
                break

            # 找到最小期望频数的组
            min_idx = np.argmin(expected)

            # 确定合并方向(与相邻的较小组合并)
            if min_idx == 0:
                merge_with = min_idx + 1  # 第一组，只能向右合并
            elif min_idx == len(expected) - 1:
                merge_with = min_idx - 1  # 最后一组，只能向左合并
            else:
                # 选择两个相邻组中较小的一个进行合并
                merge_with = min_idx - 1 if expected[min_idx - 1] < expected[min_idx + 1] else min_idx + 1

            # 执行合并(总是合并到索引较小的组)
            merge_to = min(merge_with, min_idx)
            merge_from = max(merge_with, min_idx)

            expected[merge_to] += expected[merge_from]
            observed[merge_to] += observed[merge_from]

            # 删除被合并的组
            expected = np.delete(expected, merge_from)
            observed = np.delete(observed, merge_from)
            unique = np.delete(unique, merge_from)

            # 检查是否还有足够的分组
            if len(expected) < 2:
                raise ValueError("样本量不足或分布参数不合适，无法进行有效的卡方检验")

        # 最终归一化，确保期望频数总和精确匹配观察频数总和
        expected = expected * sum(observed) / sum(expected)

        # 执行卡方检验
        statistic, p_value = chisquare(observed, expected)

        return {
            'statistic': statistic,
            'p_value': p_value,
            'df': len(observed) - 1 - 1,  # 自由度 = 组数 - 1 - 估计的参数个数
            'observed': observed,
            'expected': expected
        }

    def _get_result_message(self, test_result, alpha):
        """
        生成人类可读的检验结果描述

        参数:
            test_result: 检验结果字典
            alpha: 显著性水平

        返回:
            结果描述字符串
        """
        if 'p_value' in test_result:
            if test_result['p_value'] > alpha:
                return "不能拒绝原假设，样本分布与理论分布一致"
            else:
                return "拒绝原假设，样本分布与理论分布不一致"
        elif 'result' in test_result:
            if test_result['result']:
                return f"统计量({test_result['statistic']:.4f}) < 临界值({test_result['critical_value']:.4f})，不能拒绝原假设"
            else:
                return f"统计量({test_result['statistic']:.4f}) >= 临界值({test_result['critical_value']:.4f})，拒绝原假设"
        return ""

    def _plot_comparison(self, samples, dist_type, dist_params, test_result):
        """
        绘制样本与理论分布的比较图

        参数:
            samples: 样本数据
            dist_type: 分布类型
            dist_params: 分布参数
            test_result: 检验结果(用于标题中的统计量)
        """
        plt.figure(figsize=(12, 6))

        # 离散分布(泊松)使用条形图
        if dist_type == 'poisson':
            # 计算并绘制样本频率
            unique, counts = np.unique(samples, return_counts=True)
            plt.bar(unique, counts / len(samples), alpha=0.6,
                    label='样本频率', width=0.8)

            # 绘制理论概率质量函数(PMF)
            x = np.arange(0, max(unique) + 2)  # x轴范围
            pmf = poisson.pmf(x, dist_params['mu'])  # 计算理论PMF
            plt.plot(x, pmf, 'ro-', ms=5, label='理论PMF')  # 红色圆点连线
            plt.xticks(x)  # 设置x轴刻度为整数值
        else:
            # 连续分布使用直方图
            plt.hist(samples, bins=20, density=True,
                     alpha=0.6, label='样本直方图')

            # 绘制理论概率密度函数(PDF)
            x = np.linspace(min(samples), max(samples), 1000)  # 生成平滑的x值
            if dist_type == 'triangular':
                c = (dist_params['mode'] - dist_params['left']) / (dist_params['right'] - dist_params['left'])
                loc = dist_params['left']
                scale = dist_params['right'] - dist_params['left']
                pdf = triang.pdf(x, c=c, loc=loc, scale=scale)
            elif dist_type == 'normal':
                pdf = norm.pdf(x, loc=dist_params['mean'], scale=dist_params['std'])
            elif dist_type == 'exponential':
                pdf = expon.pdf(x, scale=dist_params['scale'])

            plt.plot(x, pdf, 'r-', lw=2, label='理论PDF')  # 红色实线

        # 设置图表标题(包含检验统计量)
        title = f'{dist_type.capitalize()}分布比较\n'
        if 'p_value' in test_result:
            title += f"检验统计量={test_result['statistic']:.4f}, p值={test_result['p_value']:.4f}"
        elif 'statistic' in test_result:
            title += f"统计量={test_result['statistic']:.4f}"
            if 'critical_value' in test_result:
                title += f", 临界值={test_result['critical_value']:.4f}"

        plt.title(title)
        plt.xlabel('值')
        plt.ylabel('概率密度/质量')
        plt.legend()  # 显示图例
        plt.grid(True)  # 显示网格
        plt.show()


# 使用示例
if __name__ == "__main__":
    # 创建分布检验器实例
    tester = DistributionTester()

    # 示例1: 三角分布检验
    print("=== 三角分布检验示例 ===")
    # 定义三角分布参数: 最小值=4, 众数=5, 最大值=8
    left, mode, right = 4, 5, 8
    # 计算形状参数c (众数在区间中的相对位置)
    # c = (mode - left) / (right - left)
    c = np.longdouble(mode - left) / np.longdouble(right - left)  # 使用高精度计算
    # 生成三角分布随机样本(200个)
    tri_samples = triang.rvs(c=c, loc=left, scale=right - left, size=20)
    # tri_samples = floor_durations
    tri_samples = std_duration
    # 执行检验(使用KS检验)
    tri_result = tester.test_distribution(
        samples=tri_samples,
        dist_type='triangular',
        dist_params={'left': left, 'mode': mode, 'right': right},
        test_type='ks'
    )
    print(f"KS统计量: {tri_result['statistic']}")
    print(f"P值: {tri_result['p_value']}")
    print(tri_result['message'])  # 打印检验结果

    # 示例2: 正态分布检验
    print("\n=== 正态分布检验示例 ===")
    # 生成标准正态分布样本(均值=0, 标准差=1)
    norm_samples = np.random.normal(loc=0, scale=1, size=200)

    # 执行检验(使用Anderson-Darling检验)
    norm_result = tester.test_distribution(
        samples=norm_samples,
        dist_type='normal',
        dist_params={'mean': 0, 'std': 1},
        test_type='anderson'
    )
    print(norm_result['message'])

    # 示例3: 泊松分布检验
    print("\n=== 泊松分布检验示例 ===")
    # 生成泊松分布样本(λ=3)
    poisson_samples = np.random.poisson(lam=3, size=200)

    # 执行检验(使用卡方检验)
    poisson_result = tester.test_distribution(
        samples=poisson_samples,
        dist_type='poisson',
        dist_params={'mu': 3},
        test_type='chi2'
    )
    print(poisson_result['message'])

    # 示例4: 指数分布检验
    print("\n=== 指数分布检验示例 ===")
    # 定义指数分布参数(scale=1/λ=2.0, 即λ=0.5)
    scale_param = 2.0
    # 生成指数分布样本
    exp_samples = np.random.exponential(scale=scale_param, size=200)

    # 执行检验(推荐使用Anderson-Darling检验)
    exp_result = tester.test_distribution(
        samples=exp_samples,
        dist_type='exponential',
        dist_params={'scale': scale_param},
        test_type='anderson'
    )
    print(exp_result['message'])